Skip to content

Conversation

Svoch
Copy link

@Svoch Svoch commented Oct 9, 2025

Add Python bindings for shard dialect. Provide means for creating constructs in this dialect in Python.

Copy link

github-actions bot commented Oct 9, 2025

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

@llvmbot llvmbot added mlir:python MLIR Python bindings mlir bazel "Peripheral" support tier build system: utils/bazel labels Oct 9, 2025
@llvmbot
Copy link
Member

llvmbot commented Oct 9, 2025

@llvm/pr-subscribers-mlir

Author: Siavash Nazari (Svoch)

Changes

Full diff: https://github.com/llvm/llvm-project/pull/162578.diff

5 Files Affected:

  • (modified) mlir/python/CMakeLists.txt (+9)
  • (added) mlir/python/mlir/dialects/ShardOps.td (+14)
  • (added) mlir/python/mlir/dialects/shard.py (+25)
  • (added) mlir/test/python/dialects/shard.py (+73)
  • (modified) utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel (+32)
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 9f5246de6bda0..20f07440df2c3 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -336,6 +336,15 @@ declare_mlir_dialect_python_bindings(
     dialects/memref.py
   DIALECT_NAME memref)
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/ShardOps.td
+  SOURCES
+    dialects/shard.py
+  DIALECT_NAME shard
+  GEN_ENUM_BINDINGS)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
diff --git a/mlir/python/mlir/dialects/ShardOps.td b/mlir/python/mlir/dialects/ShardOps.td
new file mode 100644
index 0000000000000..f8527664df67b
--- /dev/null
+++ b/mlir/python/mlir/dialects/ShardOps.td
@@ -0,0 +1,14 @@
+//===-- ShardOps.td - Entry point for ShardOps bindings ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_SHARD_OPS
+#define PYTHON_BINDINGS_SHARD_OPS
+
+include "mlir/Dialect/Shard/IR/ShardOps.td"
+
+#endif
diff --git a/mlir/python/mlir/dialects/shard.py b/mlir/python/mlir/dialects/shard.py
new file mode 100644
index 0000000000000..b64ee53bb0665
--- /dev/null
+++ b/mlir/python/mlir/dialects/shard.py
@@ -0,0 +1,25 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._shard_ops_gen import *
+from ._shard_ops_gen import _Dialect
+from ._shard_enum_gen import *
+
+try:
+    from ..ir import *
+    from ._ods_common import (
+        get_default_loc_context as _get_default_loc_context,
+        _cext as _ods_cext,
+        get_op_result_or_op_results as _get_op_result_or_op_results,
+    )
+
+    from typing import Any, List, Union
+except ImportError as e:
+    raise RuntimeError("Error loading imports from extension module") from e
+
+
+# The shard dialect currently doesn't need custom Python implementations for its operations
+# like the arith dialect does for ConstantOp. Most operations can use the generated bindings.
+# If specialized Python methods are needed for specific operations in the future,
+# they can be added here using the @_ods_cext.register_operation decorator pattern.
\ No newline at end of file
diff --git a/mlir/test/python/dialects/shard.py b/mlir/test/python/dialects/shard.py
new file mode 100644
index 0000000000000..a6455ad217fff
--- /dev/null
+++ b/mlir/test/python/dialects/shard.py
@@ -0,0 +1,73 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import shard
+from mlir.dialects import func
+
+
+def constructAndPrintInModule(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: TEST: testShardGrid
+@constructAndPrintInModule
+def testShardGrid():
+    # Test creating shard grids with different shapes
+    grid2d = shard.GridOp("grid_2d", [2, 2])
+    grid1d = shard.GridOp("grid_1d", [4])
+    grid_dynamic = shard.GridOp("grid_dynamic", [2, -1])  # -1 for dynamic dimension
+
+
+# CHECK: shard.grid @grid_2d(shape = 2x2)
+# CHECK: shard.grid @grid_1d(shape = 4)
+# CHECK: shard.grid @grid_dynamic(shape = 2x?)
+
+
+# CHECK-LABEL: TEST: testCollectiveOperations
+@constructAndPrintInModule
+def testCollectiveOperations():
+    # Create grid and types
+    grid = shard.GridOp("grid_2x2", [2, 2])
+    i32 = IntegerType.get_signless(32)
+    input_type = RankedTensorType.get([4, 2], i32)
+    gather_result_type = RankedTensorType.get([4, 4], i32)
+    
+    # Create a function to hold the operations
+    func_type = FunctionType.get([input_type], [input_type])
+    test_func = func.FuncOp("test_collectives", func_type)
+    
+    with InsertionPoint(test_func.add_entry_block()):
+        arg = test_func.entry_block.arguments[0]
+        
+        # All-gather operation
+        gather_op = shard.AllGatherOp(
+            input=arg,
+            grid=FlatSymbolRefAttr.get("grid_2x2"),
+            grid_axes=ArrayAttr.get([IntegerAttr.get(i32, 1)]),
+            gather_axis=IntegerAttr.get(i32, 1),
+            result=gather_result_type
+        )
+        
+        # All-reduce operation (ReductionKind might need different construction)
+        reduce_op = shard.AllReduceOp(
+            input=arg,
+            grid=FlatSymbolRefAttr.get("grid_2x2"),
+            reduction=IntegerAttr.get(IntegerType.get_signless(32), 1),  # 1 = sum from enum
+            result=input_type
+        )
+        
+        # Return the reduced result
+        func.ReturnOp([reduce_op])
+
+
+# CHECK: shard.grid @grid_2x2(shape = 2x2)
+# CHECK: func @test_collectives(%{{.*}}: tensor<4x2xi32>) -> tensor<4x2xi32>
+# CHECK: %{{.*}} = shard.all_gather %{{.*}} on @grid_2x2 grid_axes = [1] gather_axis = 1 : tensor<4x2xi32> -> tensor<4x4xi32>
+# CHECK: %{{.*}} = shard.all_reduce %{{.*}} on @grid_2x2 reduction = sum : tensor<4x2xi32> -> tensor<4x2xi32>
+
diff --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
index 102c4161eb74c..72af4f08bde57 100644
--- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
@@ -981,6 +981,38 @@ filegroup(
     ],
 )
 
+##---------------------------------------------------------------------------##
+# Shard dialect.
+##---------------------------------------------------------------------------##
+
+gentbl_filegroup(
+    name = "ShardOpsPyGen",
+    tbl_outs = {
+        "mlir/dialects/_shard_enum_gen.py": [
+            "-gen-python-enum-bindings",
+            "-bind-dialect=shard",
+        ],
+        "mlir/dialects/_shard_ops_gen.py": [
+            "-gen-python-op-bindings",
+            "-bind-dialect=shard",
+        ],
+    },
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "mlir/dialects/ShardOps.td",
+    deps = [
+        "//mlir:OpBaseTdFiles",
+        "//mlir:ShardTdFiles",
+    ],
+)
+
+filegroup(
+    name = "ShardOpsPyFiles",
+    srcs = [
+        "mlir/dialects/shard.py",
+        ":ShardOpsPyGen",
+    ],
+)
+
 ##---------------------------------------------------------------------------##
 # Shape dialect.
 ##---------------------------------------------------------------------------##

Copy link

github-actions bot commented Oct 9, 2025

✅ With the latest revision this PR passed the Python code formatter.

Comment on lines +988 to +1014
gentbl_filegroup(
name = "ShardOpsPyGen",
tbl_outs = {
"mlir/dialects/_shard_enum_gen.py": [
"-gen-python-enum-bindings",
"-bind-dialect=shard",
],
"mlir/dialects/_shard_ops_gen.py": [
"-gen-python-op-bindings",
"-bind-dialect=shard",
],
},
tblgen = "//mlir:mlir-tblgen",
td_file = "mlir/dialects/ShardOps.td",
deps = [
"//mlir:OpBaseTdFiles",
"//mlir:ShardTdFiles",
],
)

filegroup(
name = "ShardOpsPyFiles",
srcs = [
"mlir/dialects/shard.py",
":ShardOpsPyGen",
],
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

unless you work at Google (or you really actually do use bazel) you don't need to do this (bazel build is maintained by the users that actually use bazel)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out! I'm actually interested in building my project with Bazel, so let's keep it

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok sure np just gotta get someone that knows bazel (not me lol) to sign off. @jpienaar @superbobry this look right to you?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @makslevental! Seems like Bazel build is fine. Can I get reviews please? @jpienaar @superbobry

@Svoch Svoch force-pushed the shard-python-bindings branch from 2ea8c3e to c99fa28 Compare October 9, 2025 04:49
@Svoch Svoch changed the title Add shard Dialect Python Bindings Add shard Dialect Python Bindings Oct 9, 2025
@Svoch Svoch changed the title Add shard Dialect Python Bindings Add shard Dialect Python Bindings Oct 9, 2025
@Svoch Svoch changed the title Add shard Dialect Python Bindings [MLIR][Python] Add shard Dialect Python Bindings Oct 9, 2025
@Svoch Svoch force-pushed the shard-python-bindings branch 2 times, most recently from e1cb307 to 0dd835f Compare October 9, 2025 05:08
@Svoch
Copy link
Author

Svoch commented Oct 9, 2025

@PragmaTwice @makslevental, thanks for the reviews :) I believe I've addressed your comments. PTAL

@Svoch Svoch force-pushed the shard-python-bindings branch from 0dd835f to c8b0150 Compare October 9, 2025 05:55
@Svoch Svoch force-pushed the shard-python-bindings branch from c8b0150 to c618791 Compare October 9, 2025 05:56
@Svoch Svoch force-pushed the shard-python-bindings branch from c618791 to e3d7af0 Compare October 9, 2025 06:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bazel "Peripheral" support tier build system: utils/bazel mlir:python MLIR Python bindings mlir
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants